__author__ = 'liuyi'

import os
import sys

import getopt
import codecs
import binarize
import cPickle as pkl

from operator import itemgetter

sys.setdefaultencoding('utf-8')

def gen_word_table(strs, size = -1):
    itoa = {}
    atoi = {}
    num = 0
    if size == 0:
        return itoa, atoi
    count_table = {}

    for item in strs:
        if item in count_table:
            count_table[item] += 1
        else:
            count_table[item] = 1
    top_list = sorted(count_table.iteritems(), key=itemgetter(1), reverse=True)
    
    atoi['<UNK>'] = num
    itoa[num] = '<UNK>'
    num += 1

    if size != -1 and size != 0:
        size -= 1

    if size == -1:
        size = len(top_list)

    for item in top_list[:size]:
        if item[0] not in atoi:
            atoi[item[0]] = num
            itoa[num] = item[0]
            num += 1
    return itoa, atoi


if __name__ == '__main__':
    import sys
    import config

    try:
        opts, args = getopt.getopt(sys.argv[1:], "i:c:o:s:")
    except getopt.GetoptError:
        print "\nUsage: gen_dict.py [-i inputfile] [-c config] [-s size] [-o outputfile]\n"
        sys.exit(1)

    inFile = None
    outputFile = None
    size = -1
    for opt in opts:
        if opt[0] == '-i':
            inFile = opt[1]
        if opt[0] == '-c':
            configFile = opt[1]
        if opt[0] == '-o':
            outputFile = opt[1]
        if opt[0] == '-s':
            size = opt[1]

    if inFile == None or configFile == None or outputFile == None:
        print "\nUsage: gen_dict.py [-i inputfile] [-c config] [-s size] [-o outputfile]\n"
        sys.exit(1)

    configSet = config.CConfig(configFile)

    size = int(size)

    file = codecs.open(inFile, encoding='utf-8')
    word = []

    for line in file:
        word.extend(line.strip().split())

    itoa, atoi = gen_word_table(word,size)
    with open(configSet.pkl_dir + outputFile + '_atoi.pkl', 'wb') as f:
        pkl.dump(atoi, f, protocol = pkl.HIGHEST_PROTOCOL)

    with open(configSet.pkl_dir + outputFile + '_itoa.pkl', 'wb') as f:
        pkl.dump(itoa, f, protocol = pkl.HIGHEST_PROTOCOL)

    print outputFile + 'real size : ', len(itoa)
    print itoa
